import csv

import tensorflow as tf

from tf_record import build_int64_feature, build_float_feature, build_bytes_feature, build_example

if __name__ == '__main__':
    with open('./data/titanic_dataset.csv') as f:
        with tf.io.TFRecordWriter('./data/titanic_dataset.tfrecord') as w:
            reader = csv.reader(f, skipinitialspace=True)
            for i, record in enumerate(reader):
                if i == 0:
                    continue
                survived, pclass, name, sex, age, sibsp, parch, ticket, fare = record
                feature = {
                    'survived': build_int64_feature(int(survived)),
                    'pclass': build_int64_feature(int(pclass)),
                    'name': build_bytes_feature(str.encode(name)),
                    'sex': build_bytes_feature(str.encode(sex)),
                    'age': build_float_feature(float(age)),
                    'sibsp': build_int64_feature(int(sibsp)),
                    'parch': build_int64_feature(int(parch)),
                    'ticket': build_bytes_feature(str.encode(ticket)),
                    'fare': build_float_feature(float(fare)),
                }
                example = build_example(feature)
                w.write(example.SerializeToString())
